{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "de02cdc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import requests"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e5d60e1",
   "metadata": {},
   "source": [
    "# Document retrieval: upsert and basic query usage\n",
    "\n",
    "In this walkthrough we will go over the Retrieval API with a Azure CosmosDB Mongo vCore datastore for semantic search.\n",
    "\n",
    "Before running the notebook please initialize the retrieval API and have it running locally somewhere. Please follow the instructions to start the Retreival API provided [here](https://github.com/openai/chatgpt-retrieval-plugin#quickstart). \n",
    "\n",
    "[Azure Cosmos DB](https://azure.microsoft.com/en-us/products/cosmos-db/) Azure Cosmos DB is a fully managed NoSQL and relational database for modern app development. Using Azure Cosmos DB for MongoDB vCore, you can store vector embeddings in your documents and perform [vector similarity search](https://learn.microsoft.com/azure/cosmos-db/mongodb/vcore/vector-search) on a fully managed MongoDB-compatible database service.\n",
    "\n",
    "Learn more about Azure Cosmos DB for MongoDB vCore [here](https://learn.microsoft.com/azure/cosmos-db/mongodb/vcore/). If you don't have an Azure account, you can start setting one up [here](https://azure.microsoft.com/)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80988348",
   "metadata": {},
   "source": [
    "## Document\n",
    "\n",
    "First we will create a list of documents. From the perspective of the retrieval plugin, a [document](https://github.com/openai/chatgpt-retrieval-plugin/blob/main/models/models.py) consists of an \"id\", \"text\", \"embedding\"(optional) and a collection of \"metadata\". The \"metadata\" has \"source\", \"source_id\", \"created_at\", \"url\" and \"author\" fields. Query metadata does not expose the \"url\" field.\n",
    "\n",
    "For this example we have taken some data about a few dog breeds. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "52829ffc",
   "metadata": {},
   "outputs": [],
   "source": [
    "document_1 = {\n",
    "    \"id\": \"Siberian Husky\",\n",
    "    \"text\": \"Siberian Huskies are strikingly beautiful and energetic Arctic breed dogs known for their captivating blue eyes and remarkable endurance in cold climates.\"\n",
    "}\n",
    "\n",
    "document_2 = {\n",
    "    \"id\": \"Alaskan Malamute\",\n",
    "    \"text\": \"The Alaskan Malamute is a powerful and friendly Arctic sled dog breed known for its strength, endurance, and affectionate nature.\"\n",
    "}\n",
    "\n",
    "document_3 = {\n",
    "    \"id\": \"Samoyed\",\n",
    "    \"text\": \"The Samoyed is a cheerful and fluffy Arctic breed, renowned for its smile and gentle disposition, originally used for herding reindeer and pulling sleds in Siberia.\"\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6af96f59",
   "metadata": {},
   "source": [
    "## Indexing the Docs\n",
    "\n",
    "On the first insert, the datastore will create the collection and index if necessary on the field `embedding`. Currently hybrid search is not yet supported. \n",
    "\n",
    "To make these requests to the retrieval app API, we will need to provide authorization in the form of the BEARER_TOKEN we set earlier. We do this below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d68e796e",
   "metadata": {},
   "outputs": [],
   "source": [
    "BEARER_TOKEN_HERE = \"eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkFheXVzaCBLYXRhcmlhIiwiaWF0IjoxNTE2MjM5MDIyfQ.VHEVK_IdThXZJr8aQsfjVQ-_n4raepdpqsC5gYDsubE\"\n",
    "endpoint_url = 'http://0.0.0.0:8000'\n",
    "headers = {\n",
    "    \"Authorization\": f\"Bearer {BEARER_TOKEN_HERE}\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "954a09da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': ['doc:Siberian Husky:chunk:Siberian Husky_0',\n",
       "  'doc:Alaskan Malamute:chunk:Alaskan Malamute_0',\n",
       "  'doc:Samoyed:chunk:Samoyed_0']}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = requests.post(\n",
    "    f\"{endpoint_url}/upsert\",\n",
    "    headers=headers,\n",
    "    json={\"documents\": [document_1, document_2, document_3]\n",
    "    }\n",
    ")\n",
    "\n",
    "response.json()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "431a8616",
   "metadata": {},
   "source": [
    "## Querying the datastore\n",
    "Let's query the data store for dogs based on the place of their origin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23441d46",
   "metadata": {},
   "outputs": [],
   "source": [
    "queries = [\n",
    "    {\n",
    "        \"query\":\"I want dog breeds from Siberia.\",\n",
    "        \"top_k\":2\n",
    "    },\n",
    "    {\n",
    "        \"query\":\"I want dog breed from Alaska.\",\n",
    "        \"top_k\":1\n",
    "    }\n",
    "]\n",
    "\n",
    "response = requests.post(\n",
    "    f\"{endpoint_url}/query\",\n",
    "    headers=headers,\n",
    "    json={\"queries\":queries}\n",
    ")\n",
    "\n",
    "response.json()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "705181ee",
   "metadata": {},
   "source": [
    "## Deleting the data from the datastore\n",
    "You can either delete all the data, or provide a list of docIds to delete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b15513ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = requests.delete(\n",
    "    f\"{endpoint_url}/delete\",\n",
    "    headers=headers,\n",
    "    json={\"ids\":[\"doc:SiberianHusky:chunk:SiberianHusky_0\"]}\n",
    ")\n",
    "\n",
    "response.json()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc748e50",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = requests.delete(\n",
    "    f\"{endpoint_url}/delete\",\n",
    "    headers=headers,\n",
    "    json={\"delete_all\":True}\n",
    ")\n",
    "\n",
    "response.json()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19531965",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
